package com.itheima.datastructure.heap;

/**
 * <h3>求数组中第 K 大的元素 - 小顶堆</h3>
 * <p>
 * 思路
 * <ol>
 *     <li>向小顶堆放入前k个元素</li>
 *     <li>剩余元素</li>
 *     <ul>
 *         <li>若 <= 堆顶元素, 则略过</li>
 *         <li>若 > 堆顶元素, 则替换堆顶元素</li>
 *     </ul>
 *     <li>这样小顶堆始终保留的是到目前为止, <b>前 K 大</b>的元素</li>
 *     <li>循环结束, 堆顶元素即为<b>第 K 大</b>元素</li>
 *     <li>时间复杂度 O(n*log(n))</li>
 * </ol>
 */
public class E02Leetcode215 {

    public int findKthLargest(int[] numbers, int k) {
        MinHeap minHeap = new MinHeap(k);
        for (int i = 0; i < k; i++) {  // 向小顶椎放入前k个元素
            minHeap.offer(numbers[i]);
        }
        for (int i = k; i < numbers.length; i++) {  // 其他元素若大于堆顶元素，则替换堆顶元素
            if (numbers[i] > minHeap.peek()){
                minHeap.replace(numbers[i]);
            }
        }
        return minHeap.peek();  // 循环结束，返回堆顶元素为第k大元素。
    }

    static class MinHeap {

        int[] array;
        int size;

        public MinHeap(int capacity) {
            this.array = new int[capacity];
        }

        /**
         * 获取堆顶元素
         *
         * @return 堆顶元素
         */
        public int peek() {
            return array[0];
        }

        /**
         * 替换堆顶元素
         *
         * @param replaced 新元素
         */
        public void replace(int replaced) {
            array[0] = replaced;
            down(0);
        }

        /**
         * 堆的尾部添加元素
         *
         * @param offered 新元素
         * @return 是否添加成功
         */
        public boolean offer(int offered) {
            if (size == array.length) {
                return false;
            }
            up(offered);
            size++;
            return true;
        }

        // 将 offered 元素上浮: 直至 offered 小于父元素或到堆顶
        private void up(int offered) {
            int child = size;
            while (child > 0) {
                int parent = (child - 1) / 2;
                if (offered < array[parent]) {
                    array[child] = array[parent];
                } else {
                    break;
                }
                child = parent;
            }
            array[child] = offered;
        }


        // 将 parent 索引处的元素下潜: 与两个孩子较大者交换, 直至没孩子或孩子没它大
        private void down(int parent) {
            int left = parent * 2 + 1;
            int right = left + 1;
            int Min = parent;
            if (left < size && array[left] < array[Min]) {
                Min = left;
            }
            if (right < size && array[right] < array[Min]) {
                Min = right;
            }
            if (Min != parent) {  // 找到更大的孩子。
                swap(Min, parent);
                down(Min);
            }
        }

        // 交换两个索引处的元素
        private void swap(int i, int j) {
            int t = array[i];
            array[i] = array[j];
            array[j] = t;
        }
    }

    public static void main(String[] args) {
        E02Leetcode215 code = new E02Leetcode215();
        // 应为5
        System.out.println(code.findKthLargest(new int[]{3, 2, 1, 5, 6, 4}, 2));
        // 应为4
        System.out.println(code.findKthLargest(new int[]{3, 2, 3, 1, 2, 4, 5, 5, 6}, 4));
    }
}
